package com.android.lovely.room.chat

import android.content.Context
import android.os.Build
import android.provider.Settings
import android.telephony.TelephonyManager
import android.text.format.DateUtils
import com.android.lovely.BuildConfig
import com.android.lovely.R
import com.android.lovely.api.manager.httpLog
import com.android.lovely.const.ChatChange
import com.android.lovely.data.MessageRule
import com.android.lovely.ext.isVPNConnected
import com.android.lovely.manager.KVStore
import com.android.lovely.user.UserManager
import com.kotlin.extensions.appContext
import com.kotlin.extensions.bus.FlowBus
import com.kotlin.extensions.bus.post
import com.kotlin.extensions.json.JsonExt.fromJsonByType
import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.SharedFlow
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.sse.EventSource
import okhttp3.sse.EventSourceListener
import okhttp3.sse.EventSources
import org.json.JSONArray
import org.json.JSONObject
import java.net.URLEncoder
import java.util.TimeZone
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger

class ConversationStream {

    private var eventSource: EventSource? = null

    private val _sendFlow = MutableSharedFlow<StreamStatus>(
        replay = 1,
        extraBufferCapacity = 50,
        BufferOverflow.SUSPEND
    )
    val sendFlow: SharedFlow<StreamStatus> get() = _sendFlow

    fun request(id: Int, question: String, mode: Int, refresh: Boolean = false) {
        _sendFlow.tryEmit(StreamStatus.Start)
        
        val request = buildRequest(id, question, mode, refresh)
        val factory = createEventSourceFactory()
        
        FlowBus.post(ChatChange())
        
        eventSource = factory.newEventSource(request, createEventSourceListener())
        eventSource?.request()
    }

    private fun buildRequest(id: Int, question: String, mode: Int, refresh: Boolean): Request {
        val chatModelParam = if (UserManager.isB()) "&chat_model=$mode" else ""
        val language = URLEncoder.encode(KVStore.speakLanguage.name, "UTF-8")
        
        val url = if (refresh) {
            "${BuildConfig.API_HOST}bot/chat_stream_refresh?chat_id=${id}&lang=${language}${chatModelParam}"
        } else {
            "${BuildConfig.API_HOST}bot/chat_new?q=${
                URLEncoder.encode(question, "UTF-8")
            }&bot_id=${id}&lang=${language}&is_stream=1${chatModelParam}"
        }
        
        return Request.Builder()
            .url(url)
            .headers(createHeaders())
            .build()
    }

    private fun createHeaders() = okhttp3.Headers.Builder()
        .add("access-token", UserManager.token)
        .add("d-id", Settings.Secure.getString(appContext.contentResolver, Settings.Secure.ANDROID_ID))
        .add("version", BuildConfig.VERSION_NAME)
        .add("app-name", BuildConfig.APP_CODE)
        .add("lang", KVStore.speakLanguage.code)
        .add("sim_country", getSimCountry())
        .add("is_vpn", (if (isVPNConnected()) 1 else 0).toString())
        .add("d_model", Build.MODEL)
        .add("sys_version", Build.VERSION.RELEASE)
        .add("timezone", getTimezoneHeader())
        .build()

    private fun getSimCountry(): String {
        return (appContext.getSystemService(Context.TELEPHONY_SERVICE) as? TelephonyManager)?.simCountryIso.orEmpty()
    }


    private fun getTimezoneHeader(): String {
        val offset = TimeZone.getDefault().rawOffset / DateUtils.HOUR_IN_MILLIS
        return "GMT${if (offset >= 0) "+$offset" else "$offset"}"
    }


    private fun createEventSourceFactory(): EventSource.Factory {
        val okHttpClient = OkHttpClient.Builder()
            .retryOnConnectionFailure(false)
            .connectTimeout(3, TimeUnit.MINUTES)
            .readTimeout(3, TimeUnit.MINUTES)
            .writeTimeout(3, TimeUnit.MINUTES)
            .build()
        
        return EventSources.createFactory(okHttpClient)
    }

    private fun createEventSourceListener(): EventSourceListener {
        return StreamEventListener(_sendFlow) { 
            cancel() 
        }
    }


    fun cancel() {
        eventSource?.cancel()
    }


    sealed class StreamStatus {
        data object Start : StreamStatus()
        data class Receiving(val content: String) : StreamStatus()
        data class Failed(val code: Int, val msg: String) : StreamStatus()
        data object End : StreamStatus()
    }
}

private class StreamEventListener(
    private val statusFlow: MutableSharedFlow<ConversationStream.StreamStatus>,
    private val onCancelRequest: () -> Unit
) : EventSourceListener() {

    private var isStarted = AtomicInteger(0)
    private val analysisRule = mutableListOf<MessageRule>()
    private val isCompleted = AtomicBoolean(false)

    override fun onFailure(
        eventSource: EventSource, t: Throwable?, response: Response?
    ) {
        super.onFailure(eventSource, t, response)
        t?.printStackTrace()
        statusFlow.tryEmit(
            ConversationStream.StreamStatus.Failed(
                0, appContext.getString(R.string.send_failed)
            )
        )
    }

    override fun onClosed(eventSource: EventSource) {
        super.onClosed(eventSource)
        handleOnClose()
    }

    private fun handleOnClose() {
        if (isStarted.get() < 2) {
            statusFlow.tryEmit(
                ConversationStream.StreamStatus.Failed(
                    0, appContext.getString(R.string.send_failed)
                )
            )
            return
        }
        isStarted.set(0)
        if (!isCompleted.getAndSet(true)) {
            statusFlow.tryEmit(ConversationStream.StreamStatus.End)
        }
    }

    override fun onEvent(
        eventSource: EventSource, id: String?, type: String?, data: String
    ) {
        super.onEvent(eventSource, id, type, data)
        if (BuildConfig.DEBUG) {
            httpLog(data)
        }
        handleOnEvent(data)
    }

    private fun handleOnEvent(data: String) {
        if (isStarted.get() == 1) {
            isStarted.set(2)
        }
        if (isStarted.get() >= 1) {
            parseEventData(data, analysisRule.toMutableList(), isCompleted)
        }
        if (isStarted.get() == 0) {
            if (!parseRule(data, analysisRule)) return
            isStarted.set(1)
        }
    }


    private fun parseEventData(
        data: String, analysisRule: MutableList<MessageRule>, isCompleted: AtomicBoolean
    ) {
        runCatching {
            if (data == "[DONE]") {
                if (!isCompleted.getAndSet(true)) {
                    statusFlow.tryEmit(ConversationStream.StreamStatus.End)
                }
                return@runCatching
            }
            
            val json = JSONObject(data)
            var currentObject: Any? = json
            
            analysisRule.forEach { rule ->
                currentObject = when (rule.type) {
                    0 -> extractFromJsonArray(currentObject, rule)  // JSONArray
                    1 -> extractFromJsonObject(currentObject, rule) // JSONObject
                    2 -> extractStringValue(currentObject, rule)    // String
                    else -> currentObject
                }
            }
            
            (currentObject as? String)?.let { content ->
                statusFlow.tryEmit(ConversationStream.StreamStatus.Receiving(content))
            }
        }.onFailure {
            if (BuildConfig.DEBUG) {
                it.printStackTrace()
            }
        }
    }


    private fun extractFromJsonArray(currentObject: Any?, rule: MessageRule): Any? {
        return (currentObject as? JSONObject)?.optJSONArray(rule.key)
    }


    private fun extractFromJsonObject(currentObject: Any?, rule: MessageRule): Any? {
        return when (currentObject) {
            is JSONArray -> if (currentObject.length() > 0) {
                currentObject.optJSONObject(0)?.optJSONObject(rule.key)
            } else null
            is JSONObject -> currentObject.optJSONObject(rule.key)
            else -> null
        }
    }


    private fun extractStringValue(currentObject: Any?, rule: MessageRule): Any? {
        return when (currentObject) {
            is JSONArray -> if (currentObject.length() > 0) {
                currentObject.optJSONObject(0)?.optString(rule.key)
            } else null
            is JSONObject -> currentObject.optString(rule.key)
            else -> null
        }
    }


    private fun parseRule(data: String, analysisRule: MutableList<MessageRule>): Boolean {
        return runCatching {
            val firstJson = JSONObject(data)
            val code = firstJson.optInt("code", 0)
            
            when {
                code == 0 -> {
                    statusFlow.tryEmit(
                        ConversationStream.StreamStatus.Failed(code, firstJson.optString("detail"))
                    )
                    onCancelRequest()
                    false
                }
                code == -1 -> {
                    statusFlow.tryEmit(ConversationStream.StreamStatus.Failed(code, ""))
                    onCancelRequest()
                    false
                }
                else -> {
                    val rules = firstJson.optJSONArray("tree")
                    if (rules == null) {
                        statusFlow.tryEmit(ConversationStream.StreamStatus.Failed(code, ""))
                        onCancelRequest()
                        false
                    } else {
                        analysisRule.addAll(rules.toString().fromJsonByType<List<MessageRule>>().orEmpty())
                        true
                    }
                }
            }
        }.getOrElse { false }
    }
}
